#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#define DATA_MSG           0
#define MIN(a,b)           ((a)<(b)?(a):(b))
#define BLOCK_LOW(id,p,n)  ((id)*(n)/(p))
#define BLOCK_HIGH(id,p,n) (BLOCK_LOW((id)+1,p,n)-1)
#define BLOCK_SIZE(id,p,n) (BLOCK_HIGH(id,p,n)-BLOCK_LOW(id,p,n)+1)
#define BLOCK_OWNER(j,p,n) (((p)*((j)+1)-1)/(n))
int main(int argc, char* argv[])
{
	int rownum;//matrix row number
	int rank;//process number 
	int size;//sum of process
	int **mat;
	double starttime , endtime;
	MPI_Comm comm = MPI_COMM_WORLD;
	MPI_Init(&argc, &argv);
	MPI_Comm_size(comm, &size);
	MPI_Comm_rank(comm, &rank);
	read_row_striped_matrix(argv[1], size, rank, comm, &mat, &rownum);
	if(rank == 0)
	{
		//starttime = MPI_Wtime();
		printf("before--------------------------------------------------\n");
	}
	print_row_striped_matrix(size, rank, rownum, comm, mat);
	compute_shortest_paths(rank, size, mat, rownum, comm);
	if(rank == 0)
	{
		//endtime = MPI_Wtime();
		printf("after--------------------------------------------------\n");
	}
	print_row_striped_matrix(size, rank, rownum, comm, mat);
	//if(rank == 0)
		//printf("%lf\n", endtime - starttime);
	//free_matrix(rank, size, rownum, &mat);
	MPI_Finalize();
	return 0;
}
void read_row_striped_matrix(char* path, int size, int rank, MPI_Comm comm, int ***mat, int *rownum)
{
	int local_rownum;
	int lowrow;
	int highrow;
	int i,j;
	int *storage;
	int **lptm;
	MPI_Status status;
	FILE* file;
	if(rank == (size - 1))
	{
		file = fopen(path, "r");
		fread(rownum, sizeof(int), 1, file);
	}
	MPI_Bcast (rownum, 1, MPI_INT, size-1, comm);
	lowrow = BLOCK_LOW(rank, size, *rownum);
	highrow = BLOCK_HIGH(rank, size, *rownum);
	local_rownum = BLOCK_SIZE(rank, size, *rownum);
	storage = (int *)malloc(sizeof(int) * local_rownum * (*rownum));
	*mat = (int **)malloc(sizeof(int *) * local_rownum);
	lptm = (int *)&(*mat[0]);
	for(i=0; i<local_rownum; i++)
		(*lptm++) = &storage[i*(*rownum)];
	if(rank == (size - 1))
	{
		for(i=0; i<size-1; i++)
		{
			fread(storage, sizeof(int), BLOCK_SIZE(i, size, *rownum)*(*rownum), file);
			MPI_Send(storage, BLOCK_SIZE(i, size, *rownum)*(*rownum), MPI_INT, i, DATA_MSG, comm);
		}
		fread(storage, sizeof(int), local_rownum*(*rownum), file);
		fclose(file);
	}
	else
		MPI_Recv(storage, local_rownum*(*rownum), MPI_INT, size-1, DATA_MSG, comm, &status);
}
void print_row_striped_matrix(int size, int rank, int rownum, MPI_Comm comm, int **mat)
{
	int val=1;
	MPI_Status status;
	if(rank == 0)
	{
		print_submatrix(mat, BLOCK_SIZE(rank, size, rownum), rownum);
		MPI_Send(&val, 1, MPI_INT, rank+1, DATA_MSG, comm);
	}
	else if(rank == (size-1))
	{
		MPI_Recv(&val, 1, MPI_INT, rank-1, DATA_MSG, comm, &status);
		print_submatrix(mat, BLOCK_SIZE(rank, size, rownum), rownum);
	}
	else
	{
		MPI_Recv(&val, 1, MPI_INT, rank-1, DATA_MSG, comm, &status);
		print_submatrix(mat, BLOCK_SIZE(rank, size, rownum), rownum);
		MPI_Send(&val, 1, MPI_INT, rank+1, DATA_MSG, comm);
	}
}
void compute_shortest_paths(int id, int p, int** a, int n, MPI_Comm comm)
{
	int i,j,k;
	int offset;
	int root;
	int* tmp;
	tmp = (int*)malloc(n * sizeof(int));
	for(k=0; k<n; k++)
	{
		root = BLOCK_OWNER(k, p, n);
		if(root == id)
		{
			offset = k - BLOCK_LOW(id, p, n);
			for(j=0; j<n; j++)
				tmp[j] = a[offset][j];
		}
		MPI_Bcast(tmp, n, MPI_INT, root, comm);
		for(i=0; i<BLOCK_SIZE(id, p, n); i++)
			for(j=0; j<n; j++)
				a[i][j] = MIN(a[i][j],a[i][k]+tmp[j]);
	}
	free(tmp);
}
void print_submatrix(int **mat, int rows, int cols)
{
	int i,j;
	for(i=0; i<rows; i++)
	{
		for(j=0; j<cols; j++)
		{
			if(j == 0)
				printf("%3d", mat[i][j]);
			else
				printf(" %3d", mat[i][j]);
		}
		printf("\n");
	}
}
// void free_matrix(int rank, int size, int rownum, int ***mat)
// {
// 	int i;
// 	int local_rownum = BLOCK_SIZE(rank, size, rownum);
// 	int **lptm;
// 	lptm = (int *)&(mat[0]);
// 	for(i=0; i<local_rownum; i++)
// 		free(lptm++);
// 	free(*mat);
// }